// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/dwio/common/selective_byte_rle_column_reader.h>

namespace kumo::pollux::dwrf {

class SelectiveByteRleColumnReader
    : public dwio::common::SelectiveByteRleColumnReader {
 public:
  using ValueType = int8_t;

  SelectiveByteRleColumnReader(
      const TypePtr& requestedType,
      std::shared_ptr<const dwio::common::TypeWithId> fileType,
      DwrfParams& params,
      common::ScanSpec& scanSpec,
      bool isBool)
      : dwio::common::SelectiveByteRleColumnReader(
            requestedType,
            std::move(fileType),
            params,
            scanSpec) {
    const EncodingKey encodingKey{
        fileType_->id(), params.flatMapContext().sequence};
    auto& stripe = params.stripeStreams();
    if (isBool) {
      boolRle_ = createBooleanRleDecoder(
          stripe.getStream(
              StripeStreamsUtil::getStreamForKind(
                  stripe,
                  encodingKey,
                  proto::Stream_Kind_DATA,
                  proto::orc::Stream_Kind_DATA),
              params.streamLabels().label(),
              true),
          encodingKey);
    } else {
      byteRle_ = createByteRleDecoder(
          stripe.getStream(
              StripeStreamsUtil::getStreamForKind(
                  stripe,
                  encodingKey,
                  proto::Stream_Kind_DATA,
                  proto::orc::Stream_Kind_DATA),
              params.streamLabels().label(),
              true),
          encodingKey);
    }
  }

  void seekToRowGroup(int64_t index) override {
    dwio::common::SelectiveByteRleColumnReader::seekToRowGroup(index);
    auto positionsProvider = formatData_->seekToRowGroup(index);
    if (boolRle_) {
      boolRle_->seekToRowGroup(positionsProvider);
    } else {
      byteRle_->seekToRowGroup(positionsProvider);
    }

    POLLUX_CHECK(!positionsProvider.hasNext());
  }

  uint64_t skip(uint64_t numValues) override {
    numValues = formatData_->skipNulls(numValues);
    if (byteRle_) {
      byteRle_->skip(numValues);
    } else {
      boolRle_->skip(numValues);
    }
    return numValues;
  }

  void read(int64_t offset, const RowSet& rows, const uint64_t* incomingNulls)
      override {
    readCommon<SelectiveByteRleColumnReader, true>(offset, rows, incomingNulls);
    readOffset_ += rows.back() + 1;
  }

  template <typename ColumnVisitor>
  void readWithVisitor(const RowSet& rows, ColumnVisitor visitor);

  std::unique_ptr<ByteRleDecoder> byteRle_;
  std::unique_ptr<BooleanRleDecoder> boolRle_;
};

template <typename ColumnVisitor>
void SelectiveByteRleColumnReader::readWithVisitor(
    const RowSet& rows,
    ColumnVisitor visitor) {
  if (boolRle_) {
    if (nullsInReadRange_) {
      boolRle_->readWithVisitor<true>(
          nullsInReadRange_->as<uint64_t>(), visitor);
    } else {
      boolRle_->readWithVisitor<false>(nullptr, visitor);
    }
  } else {
    if (nullsInReadRange_) {
      byteRle_->readWithVisitor<true>(
          nullsInReadRange_->as<uint64_t>(), visitor);
    } else {
      byteRle_->readWithVisitor<false>(nullptr, visitor);
    }
  }
}

} // namespace kumo::pollux::dwrf
